### This is a Wrapper for imputation functions to be used in other scripts. ###
# Impute_default_all is used in analyzing real and splatter simulated data.
# Individual wrappers are used for testing effects of method parameters.

# Format required to add additional imputation methods:
#	input: a SingleCellExperiment object with "counts" and "logcounts"
#		a parameter value (with default), 
#		whether the input should be normalized (with default=TRUE)
#		number of cores to use (if applicable)
#		random seed to ensure reproducibility (with default).
#	output: the imputed matrix where rows=genes, cols=cells

require("SingleCellExperiment")
require("Rmagic") 
require("DrImpute") 
require("scImpute") 
require("SAVER") 
require("doMC")
require("tictoc")
source("../knn-smoothing/knn_smooth.R") 
source('./Gmm/odgmm.R')
source('./Gmm/gmm_impute.R') 


scImpute_wrapper<- function(sce, param=0.5, do.norm=TRUE, n.cores=16, seed=42) {
    # normalization is never used
    tmp_dir <- paste(paste("Tmp",round(runif(1)*100000), sep="_"),"/",sep="")
    dir.create(tmp_dir, showWarnings = FALSE)
    set.seed(seed)
    require(doParallel)
    n.cores <- as.numeric(n.cores)
    cl <- makeCluster(n.cores)
    registerDoParallel(cl)
    param_name <- "Dropout Threshold"
    registerDoParallel(cores = n.cores)
    saveRDS(as.matrix(assays(sce)[["counts"]]), file=paste(tmp_dir,"tmp.rds", sep=""))

    print(dim(assays(sce)[["counts"]]))

    scImpute::scimpute(
        paste(tmp_dir, "tmp.rds", sep="/"), 
        infile="rds", 
        outfile="rds",
        type="count", 
        drop_thre=param, 
        out_dir=tmp_dir,
        Kcluster=ifelse(length(unique(sce$Group)) == 0, 1, length(unique(sce$Group))),
        ncores=n.cores
    )
    out <- readRDS(paste(tmp_dir,"scimpute_count.rds", sep=""))
    stopCluster(cl)
    unlink(tmp_dir, recursive=TRUE)
    return(out);
}

DrImpute_wrapper <- function(sce, param=0, do.norm=TRUE, seed=42) {
	# uses pre-defined log-normalized matrix
	set.seed(seed)
	param_name <- "Zeros Remaining"
	out <- DrImpute::DrImpute(
      as.matrix(assays(sce)[["logcounts"]]),
      ks=ifelse(length(unique(sce$Group)) == 0, 10, length(unique(sce$Group))),
      zerop=param
  )
	return(out)
}

SAVER_wrapper <- function(sce, param=1, do.norm=TRUE, seed=42, n.cores=16){
	# optional CPM-like normalization
	n.cores <- as.numeric(n.cores)
	set.seed(seed)
	require(doParallel)
	cl <- makeCluster(n.cores)
	registerDoParallel(cl)
	param_name <- "Percent of Genes"
	sf <- 1
	if (do.norm) {
		sf <- NULL
	}
	if (param == 1) {
      out <- saver(
          assays(sce)[["counts"]], 
          do.fast=TRUE, 
          size.factor=sf, 
          ncores=n.cores
      )
  } else {
      out <- saver(
          assays(sce)[["counts"]], 
          do.fast=TRUE, 
          size.factor=sf,
          npred=nrow(sce)*param, 
          ncores=n.cores
      )
  }
	stopCluster(cl)
	return(out$estimate);
}

MAGIC_k_wrapper <- function(sce, param=10, do.norm=TRUE, seed=42, n.jobs=10) {
	# optional CPM-like normalization
	set.seed(seed)
	param_name <- "K neighbours"
	out <- magic(
    as.matrix(t(assays(sce)[["counts"]])), 
		k=param,
    seed=seed,
    n.jobs = n.jobs
  )
	return(t(out))
}

MAGIC_wrapper <- function(sce, param=10, seed=42, n.jobs=10) {
	# optional CPM-like normalization
	param_name <- "K neighbours"
	out <- magic(
    as.matrix(t(assays(sce)[["counts"]])), 
		k=param,
    seed=seed,
    n.jobs = n.jobs
  )
	return(t(as.data.frame(out)))
}

knn_wrapper <- function(sce, param=ncol(sce)/20, do.norm=TRUE, seed=42) {
	# includes CPM-like normalization
	param_name <- "K neighbours"
	out <- knn_smoothing(as.matrix(assays(sce)[["counts"]]), k=param, d=10, seed=seed)
  return(out)
}



## First step:
##after loading GMM model's all functions
##fitting for single cell matrix: row is genes column is cell
GMM_impute_wrapper <- function(
    sim_sce, 
    seed = 42, 
    comp = 5, 
    n.cores = 10 , 
    rate = 0.5, 
    param=10
) {
      
    mu0 = 0
    sigma0 = 0.1
    set.seed(seed)
    comp = comp

    # gene * cell
    X_tab = as.matrix(assays(sim_sce)[["counts"]])
    
    registerDoMC(n.cores)
    
    # remove redundant variables
    tic("Filter")
    res_multi <- foreach(i=1:dim(X_tab)[1]) %dopar% {
        X = as.numeric(X_tab[i,])   #X_tab is the matrix need to be fitted by GMM
        
        res = tryCatch(
            odgmm(X,comp), 
            error=function(e) {
                return(NA)
            }
        )
                
        if (length(res) == 0 || sum(is.na(unlist(res))) != 0) {
            return(NULL)
        } else {
            return(proc_res(res)) 
        }
    }
    toc()
    
    # cell * gene
    X = as.data.frame(matrix(NA, nrow = dim(X_tab)[2], ncol = dim(X_tab)[1]))
    
    colnames(X) = rownames(X_tab)
    rownames(X) = colnames(X_tab)
    
    xvals = list()
    
    for (i in 1:length(res_multi)) {
        temp_res = res_multi[[i]]
        if(!is.null(temp_res)) {
            # print(head(temp_res))
            X[, i] = as.numeric(temp_res[[1]])
            xvals[[length(xvals) + 1]] = temp_res[[2]]
        } 
    }
    
    # remove na
    X = as.matrix(t(na.omit(as.matrix(t(X)))))
    
    n_gene = dim(X_tab)[1]
    n_cell = dim(X_tab)[2]
        
    keep_rate = apply(X > 0, 2, sum) / n_cell
    gene_inds = which(keep_rate > rate)

    if(length(gene_inds) == 0) {
        stop("rate is too high")
    }
    
    print(summary(keep_rate))
    
    # perform imputation
    XX = X[, gene_inds]
    
    XX_vals = xvals[gene_inds]
    res_after_impute = gmm_impute(
        orig_data=as.matrix(t(X_tab[gene_inds, ])), 
        gmm_label_mat=as.matrix(XX), 
        gmm_mus_list=XX_vals,
        k=param
    )
    
    # X_impute = res_after_impute$impute_label_mat

    res = as.matrix(t(res_after_impute$impute_val_mat))
    print(dim(res))
    new_res = X_tab

    new_res = matrix(0, nrow=nrow(X_tab), ncol=ncol(X_tab))
    rownames(new_res) = rownames(X_tab)
    colnames(new_res) = colnames(X_tab)


    for(i in rownames(res)) {
        new_res[i, colnames(new_res)] = res[i, ]
    }

    return(new_res)
}


GMM_impute_wrapper_v1 <- function(
    sim_sce, 
    seed = 42, 
    comp = 5, 
    n.cores = 10 , 
    rate = 0.8, 
    param=NULL
) {
    mu0 = 0
    sigma0 = 0.1
    set.seed(seed)
    comp = comp
    X_tab = log2(as.matrix(assays(sim_sce)[["counts"]]) + 1)
    
    registerDoMC(n.cores)
    
    # remove redundant variables
    tic("Filter")
    res_multi <- foreach(i=1:dim(X_tab)[1]) %dopar% {
        X = as.numeric(X_tab[i,])   #X_tab is the matrix need to be fitted by GMM
        
        res = tryCatch(
            odgmm(X,comp), 
            error=function(e) {
                return(NA)
            }
        )
                
        if (length(res) == 0 || sum(is.na(unlist(res))) != 0) {
            return(NULL)
        } else {
            return(proc_res(res)) 
        }
    }
    toc()
    
    X = as.data.frame(matrix(NA, nrow = dim(X_tab)[2], ncol = dim(X_tab)[1]))
    
    colnames(X) = rownames(X_tab)
    rownames(X) = colnames(X_tab)
    
    xvals = list()
    
    for (i in 1:length(res_multi)) {
        temp_res = res_multi[[i]]
        if(!is.null(temp_res)) {
            # print(head(temp_res))
            X[, i] = as.numeric(temp_res[[1]])
            xvals[[length(xvals) + 1]] = temp_res[[2]]
        } 
    }
    
    X = as.matrix(t(na.omit(as.matrix(t(X)))))
    
    n_gene = dim(X_tab)[1]
    n_cell = dim(X_tab)[2]
        
    keep_rate = apply(X > 0, 2, sum) / n_cell
    gene_inds = which(keep_rate > rate)

    if(length(gene_inds) == 0) {
        stop("rate is too high")
    }
    
    print(summary(keep_rate))
    
    # perform imputation
    XX = X[, gene_inds]
    
    XX_vals = xvals[gene_inds]
    res_after_impute = gmm_impute(
        orig_data=as.matrix(t(X_tab[gene_inds, ])), 
        gmm_label_mat=as.matrix(XX), 
        gmm_mus_list=XX_vals,
        k=param
    )
    
    res = as.matrix(t(res_after_impute$impute_val_mat))

    res = rbind(res, X_tab[!rownames(X_tab) %in% rownames(res), ])
    
    res = res[rownames(X_tab), colnames(X_tab)]
    return(res)
}


Impute_default_all <- function(sce, n.cores=16) {
    res <- scImpute_wrapper(sce, n.cores=n.cores)
    assays(sce)[["sci"]] <- res;
    
    res <- DrImpute_wrapper(sce)
    assays(sce)[["dri"]] <- res;
    
    res <- MAGIC_wrapper(sce)
    assays(sce)[["magic"]] <- res;
    
    res <- knn_wrapper(sce)
    assays(sce)[["knn"]] <- res;
    
    res <- SAVER_wrapper(sce, n.cores=n.cores)
    assays(sce)[["saver"]] <- res;
    
    res <- GMM_impute_wrapper(sce, rate = 0.5)
    assays(sce)[["old_GMM"]] <- res
    
    res <- GMM_impute_wrapper1(sce, rate = 0.5)
    assays(sce)[["new_GMM"]] <- res
    
    return(sce)
}

